package com.example.performance.r2dbc.compare;

import com.example.performance.r2dbc.entity.RecordRowMapper;
import com.example.performance.r2dbc.utils.PerformanceCollector;
import com.example.performance.r2dbc.utils.TestSuite;
import io.r2dbc.postgresql.PostgresqlConnectionConfiguration;
import io.r2dbc.postgresql.PostgresqlConnectionFactory;
import io.r2dbc.postgresql.util.LogLevel;
import io.r2dbc.spi.Connection;
import io.r2dbc.spi.ConnectionFactory;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.reactivestreams.Publisher;
import org.springframework.stereotype.Service;
import org.springframework.util.StopWatch;
import reactor.core.publisher.Flux;
import reactor.core.scheduler.Schedulers;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.*;
import java.util.function.Function;
import java.util.logging.Level;

/**
 * @author zpMeng
 * @since 6/19/2025
 */
@Slf4j
@Service
@RequiredArgsConstructor
public class R2DBCService {

    private final ConnectionFactory connectionFactory;

    public static final ConnectionFactory PG_CONNECTION_FACTORY = new PostgresqlConnectionFactory(
            PostgresqlConnectionConfiguration.builder()
                    .host("10.86.0.7")
                    .port(5432)
                    .username("fmuser_atlas_demo")
                    .password("fm")
                    .database("fmuser_atlas_demo")
                    .autodetectExtensions(true)
                    .errorResponseLogLevel(LogLevel.DEBUG)
                    .noticeLogLevel(LogLevel.DEBUG)
                    .build());

    public void testSerialPerformance() {
        Map<Long, Integer> nodeSizeToCounts = TestSuite.prepareSerialTestSuite();
        System.out.println(nodeSizeToCounts);

        for (Map.Entry<Long, Integer> nodeSizeToCount : nodeSizeToCounts.entrySet()) {

            Long nodeSize = nodeSizeToCount.getKey();
            Integer testTimes = nodeSizeToCount.getValue();

            ArrayList<Double> durations = new ArrayList<>();
            for (int i = 0; i < testTimes; i++) {
                double duration = executeOneR2DBCQuery(0, nodeSize);
                durations.add(duration);
            }

            PerformanceCollector.printSerial(Map.of(nodeSize, durations));
        }
    }

    public void testParallelPerformance(int start, int end, int nodeSize) throws Exception {
        Map<Integer, Integer> concurrencyToTestTimes = TestSuite.prepareParallelTestSuite(start, end, nodeSize);
        System.out.println(concurrencyToTestTimes);

        for (Map.Entry<Integer, Integer> entry : concurrencyToTestTimes.entrySet()) {
            Integer concurrency = entry.getKey();
            Integer count = entry.getValue();

            executeOneParallelTest(count, concurrency, nodeSize);
        }
    }


    public void executeOneParallelTest(int count, int concurrency, int nodeSize) throws InterruptedException {
        List<Double> totalDurations = new ArrayList<>();

        try (ExecutorService executor = Executors.newFixedThreadPool(concurrency)) {
            for (int i = 1; i <= count; i++) {
                testOneConcurrencyPerformance(executor, concurrency, nodeSize, totalDurations);
            }
        }

        PerformanceCollector.printParallel(totalDurations, concurrency);
    }

    public void testOneConcurrencyPerformance(ExecutorService executorService, int concurrency, long nodeSize,
                                              List<Double> collectors) throws InterruptedException {
        CountDownLatch latch = new CountDownLatch(concurrency);
        List<Future<Double>> futures = new ArrayList<>(concurrency);
        StopWatch stopWatch = new StopWatch();
        stopWatch.start();

        for (int i = 0; i < concurrency; i++) {
            int version = getVersion(getVersion(concurrency));

            Future<Double> currentTask = executorService.submit(() -> {
                try {
                    return executeOneR2DBCQuery(version, nodeSize);
                } finally {
                    latch.countDown();
                }
            });
            futures.add(currentTask);
        }

        latch.await();
        stopWatch.stop();

        // 收集每个任务的执行时间
        for (Future<Double> future : futures) {
            try {
                collectors.add(future.get());
            } catch (ExecutionException e) {
                System.err.println("Task failed: " + e.getCause());
            }
        }
    }

    private static int getVersion(int concurrency) {
        return concurrency % 20;
    }


    public double executeOneR2DBCQuery(int version, Long nodeSize) {
        String sql = "select * from DS_AGG_427_%d where node_id <= %d".formatted(version, nodeSize);

        StopWatch stopWatch = new StopWatch();
        stopWatch.start("nodeSize: " + nodeSize);

        RecordRowMapper recordRowMapper = new RecordRowMapper();

        Flux
                .usingWhen(
                        connectionFactory.create(),
                        (Function<Connection, Publisher<?>>) connection -> {
                            // log.info("AutoCommit: {}, Thread: {}", connection.isAutoCommit(), Thread.currentThread().getName());
                            return Flux
                                    .from(connection
                                            .createStatement(sql)
                                            //.fetchSize(1000000)
                                            .execute()
                                    )
                                    .flatMap(result -> result.map(recordRowMapper))
                                    .count()
                                    .subscribeOn(Schedulers.boundedElastic());
                        },
                        Connection::close
                )
                .log("reactor test", Level.ALL)
                .blockLast();

        stopWatch.stop();

        return stopWatch.getTotalTimeSeconds();
    }

}
